package edu.utdallas.cs.securechat.main;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.net.InetAddress;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;
import java.util.Arrays;
import java.util.Date;
import java.util.HashMap;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;

public class ClientProtocol {
	SecretKey clientKey;
	SecretKey authenticationKey;
	SecretKey sessionKey;
	Key	serverPublicKey;
	Key clientPrivateKey;
	Key clientPublicKey;
	byte[] TicketGrantingTicket;
	String username;
	long ticketExpiration;
	final int keysize = 16;
	
	HashMap<String, ClientSession> userSessions;
	
	public ClientProtocol(){
		this.userSessions = new HashMap<String, ClientSession>();
		try {
			this.generateRSAKeys();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}
	
	public ClientProtocol(String username)
	{
		this.username = username;
		this.userSessions = new HashMap<String, ClientSession>();
		try{
			this.generateRSAKeys();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}
	
	public void setUsername(String username){
		this.username = username;
	}
	
	public String getUsername(){
		return this.username;
	}
	
	public SecretKey getClientKey()
	{
		return clientKey;
	}

	public void setSecretKey(SecretKey key)
	{
		this.clientKey = key;
	}
	
	public void setUserAddress(String username, InetAddress address){
		this.userSessions.get(username).setAddress(address);
	}
	
	public Key getPublicKey(){
		return this.clientPublicKey;
	}
	
	public void setSessionKey(SecretKey key)
	{
		this.sessionKey = key;
	}
	
	public long getTicketExpiration()
	{
		return this.ticketExpiration;
	}
	
	public String hashPassword(byte[] pass) throws NoSuchAlgorithmException{
		MessageDigest hash = MessageDigest.getInstance("SHA-1");
		hash.reset();
	    hash.update(pass);
	    //hash.digest();
        String s = new BigInteger(1, hash.digest()).toString(16);
        return (s.length() % 2 == 0) ? s : "0" + s;
	}
	
	public byte[] requestTicket(String user, long timestamp) throws IOException, InvalidKeyException, NoSuchAlgorithmException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException{
		ByteArrayOutputStream request = new ByteArrayOutputStream();  
		DataOutputStream dos = new DataOutputStream(request);  
		dos.writeLong(timestamp);  
		dos.flush();  
		request.write(user.getBytes());
		
		ByteArrayOutputStream result = new ByteArrayOutputStream();
		byte[] Erequest = this.AESEncrypt(request.toByteArray(), this.sessionKey);
		result.write(Erequest.length);
		result.write(Erequest);
		result.write(this.TicketGrantingTicket);
			
		return result.toByteArray();
	}
	
	public void handleTicketResponse(byte[] response, long timestamp) throws Exception{
		ByteArrayInputStream responseStream = new ByteArrayInputStream(response);
		int l = (int)responseStream.read();
		byte[] auth = new byte[l];
		byte[] Esession = new byte[responseStream.available()-auth.length];
		
		responseStream.read(auth);
		responseStream.read(Esession);
		
		ByteArrayInputStream sessionStream = new ByteArrayInputStream(this.AESDecrypt(Esession, this.sessionKey));
		
		DataInputStream dis = new DataInputStream(sessionStream);
		long plus1 = dis.readLong();
		long expiration = dis.readLong();

		byte[] k = new byte[keysize];
		byte[] u = new byte[sessionStream.available()-k.length];

		sessionStream.read(k);
		sessionStream.read(u);
		
		if (plus1 == timestamp+1){
		
			String username = new String(u);
			SecretKey sessionKey = new SecretKeySpec(k, "AES");

			if (this.userSessions.containsKey(username)){
				userSessions.remove(username);
			}
			userSessions.put(username, new ClientSession(username, auth, sessionKey, expiration));
		} else {throw new Exception("Server timestamp is incorrect: possible reply attack");}
	}
	
	public boolean isConnected(String user){
		if (this.userSessions.containsKey(user)){
			ClientSession session = this.userSessions.get(user);
			return (session.isConfirmed() && (new Date().getTime()<=session.getExpiration()));
		} else { return false; }
		
	}
	
	public byte[] encryptMessage(byte[] message, String user) throws Exception{
		if (this.isConnected(user)){
			return this.AESEncrypt(message, this.userSessions.get(user).getKey());
		} else {throw new Exception("Not connected to user:" + user);}
	}
	
	public byte[] decryptMessage(byte[] message, String user) throws Exception{
		if (this.isConnected(user)){
			return this.AESDecrypt(message, this.userSessions.get(user).getKey());
		} else {throw new Exception("Not connected to user:" + user);}
	}
	
	public byte[] authenticateToClient(String name, long timestamp) throws IOException, InvalidKeyException, NoSuchAlgorithmException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException{
		ClientSession session = this.userSessions.get(name);
		
		ByteArrayOutputStream requestStream = new ByteArrayOutputStream();  
		DataOutputStream dos = new DataOutputStream(requestStream);  
		dos.writeLong(timestamp);  	
		dos.flush();
				
		requestStream.write(this.username.getBytes());
				
		ByteArrayOutputStream resultStream = new ByteArrayOutputStream();
		byte[] Erequest = this.AESEncrypt(requestStream.toByteArray(), session.getKey());
		resultStream.write(Erequest.length);
		resultStream.write(Erequest);
		resultStream.write(session.getAuth());
		
		return resultStream.toByteArray();
	}
	
	public byte[] checkUserAuthentication(byte[] message, Key publicKey) throws Exception{
		ByteArrayInputStream messageStream = new ByteArrayInputStream(message);
		int l = messageStream.read();
		byte[] Erequest = new byte[l];
		messageStream.read(Erequest);
		byte[] Eauth = new byte[messageStream.available()];
		messageStream.read(Eauth);
		
		ByteArrayInputStream authStream = new ByteArrayInputStream(this.AESDecrypt(Eauth, this.clientKey));
		DataInputStream dis = new DataInputStream(authStream);
		long expire = dis.readLong();
		byte[] k = new byte[keysize];
		authStream.read(k);
		byte[] au = new byte[authStream.available()];
		authStream.read(au);
		String authUser = new String(au);
		
		if (expire >= new Date().getTime()){
			SecretKey key = new SecretKeySpec(k, "AES");
			ByteArrayInputStream requestStream = new ByteArrayInputStream(this.AESDecrypt(Erequest, key));
			dis = new DataInputStream(requestStream);
			long timestamp = dis.readLong();
			byte[] ru = new byte[requestStream.available()];
			requestStream.read(ru);
			String requestUser = new String(ru);
			
			if (requestUser.equals(authUser)){
				if (this.userSessions.containsKey(requestUser)){userSessions.remove(requestUser);}
				ClientSession session = new ClientSession(requestUser, key, publicKey, expire, true);
				this.userSessions.put(requestUser, session);
				ByteArrayOutputStream resultStream = new ByteArrayOutputStream();
				DataOutputStream dos = new DataOutputStream(resultStream);
				dos.writeLong(timestamp+1);
				resultStream.write(this.username.getBytes());
				return this.AESEncrypt(resultStream.toByteArray(), key);
				
			} else {throw new Exception("Username in request does not match ticket");}
			
			
		} else {throw new Exception("Ticket is Expired. Request a new ticket from Server");}
		
	}
	
	public void confirmUserConnection(byte[] message, String user, Key publicKey, long timestamp) throws Exception{
		if (this.userSessions.containsKey(user)){
			ByteArrayInputStream messageStream = new ByteArrayInputStream(this.AESDecrypt(message, this.userSessions.get(user).getKey()));
			DataInputStream dis = new DataInputStream(messageStream);
			long tplus1 = dis.readLong();
			byte[] u = new byte[messageStream.available()];
			messageStream.read(u);
			if (user.equals(new String(u)) && (tplus1==(timestamp+1))){
				this.userSessions.get(user).confirm();
				this.userSessions.get(user).setPublicKey(publicKey);
			} else {throw new Exception("Authentication failed");}
			                    
		} else {throw new Exception("Invalid user");}
	}
	
	public void HandleTGT(byte[] TGT) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, IOException
	{
		byte[] response = new byte[keysize+Long.SIZE/8+8];
		byte[] ticket = new byte[TGT.length-response.length];
		
		ByteArrayInputStream TGTstream = new ByteArrayInputStream(TGT);
		TGTstream.read(response);
		TGTstream.read(ticket);
		this.TicketGrantingTicket = ticket;
		
		ByteArrayInputStream responseStream = new ByteArrayInputStream(this.AESDecrypt(response, clientKey));
		DataInputStream dis = new DataInputStream(responseStream);
		long expire = dis.readLong();
		byte[] sKey = new byte[keysize];
		
		responseStream.read(sKey);
		this.sessionKey = new SecretKeySpec(sKey, "AES");
		this.ticketExpiration = expire;
	}
	
	public byte[] hashMessage(byte[] message) throws NoSuchAlgorithmException
	{
		MessageDigest hash = MessageDigest.getInstance("SHA-1");
		hash.reset();
	    hash.update(message);
	    return hash.digest();
	}
	
	public byte[] RSADecrypt(byte[] message) throws NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException
	{
	    Cipher cipher = Cipher.getInstance("RSA");
	    cipher.init(Cipher.DECRYPT_MODE, this.clientPrivateKey);
	    byte[] plainText = cipher.doFinal(message);
	    return plainText;
	}
	
	public byte[] RSADecrypt(byte[] message, Key key) throws NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException
	{
	    Cipher cipher = Cipher.getInstance("RSA");
	    cipher.init(Cipher.DECRYPT_MODE, key);
	    byte[] plainText = cipher.doFinal(message);
	    return plainText;
	}
	
	public byte[] signMessage(byte[] message) throws NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException, IOException
	{
		byte[] hash = this.RSAEncrypt(hashMessage(message));
	    ByteArrayOutputStream signedMessage = new ByteArrayOutputStream();
	    signedMessage.write(hash);
	    signedMessage.write(message);
	    return signedMessage.toByteArray();
	}
	
	public byte[] checkMessage(byte[] signedMessage, String username) throws Exception{
		Key userPublicKey = this.userSessions.get(username).getPublicKey();
		ByteArrayInputStream m = new ByteArrayInputStream(signedMessage);
		byte[] hash = new byte[64];
		m.read(hash);
		byte[] message = new byte[m.available()];
		m.read(message);
		if (Arrays.equals(this.RSADecrypt(hash, userPublicKey), this.hashMessage(message))){
			return message;
		} else {
			throw new Exception("Message failed integrity check");
		}
		
		
	}
	
	public byte[] RSAEncrypt(byte[] message) throws NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException
	{
		Cipher cipher = Cipher.getInstance("RSA");
	    cipher.init(Cipher.ENCRYPT_MODE, this.clientPrivateKey);
	    byte[] cipherText = cipher.doFinal(message);
	    return cipherText;
	}
	
	public byte[] RSAEncrypt(byte[] message, Key publicKey) throws NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException
	{
		Cipher cipher = Cipher.getInstance("RSA");
	    cipher.init(Cipher.ENCRYPT_MODE, publicKey);
	    byte[] cipherText = cipher.doFinal(message);
	    return cipherText;
	}
	
	public byte[] AESEncrypt(byte[] message, SecretKey key) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException
	{
		Cipher c = Cipher.getInstance("AES");
		SecretKeySpec k = new SecretKeySpec(key.getEncoded(), "AES");
		c.init(Cipher.ENCRYPT_MODE, k);
		byte[] encryptedData = c.doFinal(message);
		return encryptedData;
	}
	
	public byte[] AESDecrypt(byte[] message, SecretKey key) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException
	{
		Cipher c = Cipher.getInstance("AES");
		SecretKeySpec k = new SecretKeySpec(key.getEncoded(), "AES");
		c.init(Cipher.DECRYPT_MODE, k);
		byte[] data = c.doFinal(message);
		return data;
	}
	public void keyFromPassword(char[] pass) throws NoSuchAlgorithmException, InvalidKeySpecException
	{
		byte[] bytes = new byte[pass.length * 2];
		for (int i = 0; i < pass.length; i++) {
			bytes[i * 2] = (byte) (pass[i] >> 8);
			bytes[i * 2 + 1] = (byte) pass[i];
		}
		SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1");
		KeySpec spec = new PBEKeySpec(this.username.concat(this.hashPassword(bytes)).toCharArray(), new byte[]{4, 8, 15, 16, 23, 42}, 65536, 128);
		SecretKey tmp = factory.generateSecret(spec);
		SecretKey secret = new SecretKeySpec(tmp.getEncoded(), "AES");
		this.clientKey = secret;
	}
	
	public void generateRSAKeys() throws NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException
	{
	    SecureRandom random = new SecureRandom();
	    KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");

	    generator.initialize(512, random);

	    KeyPair pair = generator.generateKeyPair();
	    this.clientPublicKey = pair.getPublic();
	    this.clientPrivateKey = pair.getPrivate();
	}
}